import torch

from ignite.metrics import Metric

from ignite.exceptions import NotComputableError
from ignite.metrics.metric import sync_all_reduce, reinit__is_reduced

class L2SD(Metric):


    def __init__(self, output_transform=lambda x: x, device="cpu"):
        self._samples = None
        super(L2SD, self).__init__(output_transform=output_transform, device=device)


    @reinit__is_reduced
    def reset(self):
        self._samples = torch.empty((0,), device=self._device)
        # super(L2SD, self).reset()
        super().reset()


    @reinit__is_reduced
    def update(self, output):
        if output[0] is None:
            return
        l2 = output[0].detach().to(self._device)
        self._samples = torch.cat((self._samples, l2), 0)


    @sync_all_reduce("_num_examples", "_num_correct:SUM")
    def compute(self):
        if self._samples.size(0) == 0:
            raise NotComputableError("L2SD has 0 sample!")
        res = torch.std(self._samples)
        return res


if __name__ == "__main__":
    m = L2SD()

    dis = torch.randint(10, (4, 10)).float()
    print(dis)

    for i in range(dis.size(0)):
        m.update((dis[i], None))
    
    res = m.compute()

    label = torch.std(torch.flatten(dis))
    print(label)



